#include <string.h>

#include "lwip/stats.h"
#include "lwip/mem.h"
#include "lwip/udp.h"
#include "lwip/ip_addr.h"
#include "lwip/netif.h"
#include "lwip/inet.h"
#include "netif/etharp.h"

#include "lwip/sys.h"
#include "lwip/opt.h"
#include "lwip/dns.h"

#if LWIP_DNS

#define DNS_QUERY_TIMEOUT        300 // ticks

static u16_t id = 0x3322;
struct ip_addr answer;
static int volatile flag = 0;  // set to 1 is server return a valid reply, -1 on error
void dns_recv(void *arg, struct udp_pcb *pcb, struct pbuf *p, struct ip_addr *addr, u16_t port);




err_t dns_start(struct netif *netif, char *host, struct ip_addr *addr)
{

        struct pbuf *p;
        struct udp_pcb *pcb; 
        unsigned int jiffie, i, j;
        struct dns_hdr *dns_header;
        unsigned char qname[MAX_DOMAIN_NAME_LEN + 3];  // two for count, additional 1 byte might used to count length if MAX_DOMAIN_NAME_LEN is input
        unsigned char c[4] = {0, 1, 0, 1}; // type & class
        err_t err = ERR_CONN;
                
        if(netif == NULL || host == NULL || addr == NULL) 
                return(ERR_ARG);        	        	
        
        if(strlen(host) >= MAX_DOMAIN_NAME_LEN)
        	return(ERR_ARG);  // you sure this is the correct host name?
        
        IP4_ADDR(&answer, 0, 0, 0, 0);
        if(ip_addr_cmp(&(netif->dns), &answer))
                return(ERR_RTE);        // DNS server not set. cannot route to no where        
        
        // prepare qurey name
        memcpy(qname + 1, host, strlen(host));
        qname[strlen(host) + 2]  = qname[strlen(host) + 1]= 0;
        for(i = 1; i < strlen(host) + 1; i++) {
                if(qname[i] == '.')
        	        qname[i] = 0x00;        	
        }
        i = 1;
        while((j = strlen((char *)&qname[i])) != 0) {
                if(j >= 63)
                	return(ERR_ARG);
                
                qname[i - 1] = j;
                i += (j + 1); // j bytes characters and 1 byte count                
                if(j == 0)
                        break;
        
        }
        if((pcb = udp_new())== NULL)
                return(ERR_MEM);        
        
        if((p = pbuf_alloc(PBUF_TRANSPORT, strlen(host) + 2 + 4 + sizeof(struct dns_hdr), PBUF_RAM)) == NULL) { // 4: type & classs, 2: first and last count.
                udp_remove(pcb);
                return(ERR_MEM);
        }        
        	
        
        // prepare DNS packet
        dns_header = p->payload;
        //memset(*dns_header, 0, sizeof(dns_header));
        dns_header->id = htons(id);
        dns_header->flags = htons(DNS_QUERY|DNS_SQUERY|DNS_RD);
        dns_header->tot_que = htons(1);
        dns_header->tot_ans = dns_header->tot_auth = dns_header->tot_add = 0;
        memcpy((char *)p->payload + sizeof(struct dns_hdr), qname, strlen(host) + 2);  // query name
        memcpy((char *)p->payload + sizeof(struct dns_hdr) + strlen(host) + 2, c, 4);  // type & class
  
        flag = 0;
        
        udp_bind(pcb, IP_ADDR_ANY, 0);      // get us a fixed port
       // udp_connect(pcb, &netif->dns, DNS_SERVER_PORT);
        udp_recv(pcb, dns_recv, 0);        
        udp_sendto(pcb, p, &netif->dns, DNS_SERVER_PORT);
        // polling for 1 second to see if answer become not 0
        jiffie = sysGetTicks(TIMER0);
        while(sysGetTicks(TIMER0) - jiffie < DNS_QUERY_TIMEOUT) {
        	// server replied?
                if(flag == 1) {
                        *addr = answer;
                        err = ERR_OK;
                        break;
                }
                if(flag == -1)
                        break;
        }
        pbuf_free(p);
        udp_remove(pcb);
        id++;
        return(err);
}	

// extract first RR from server reply and return
void dns_recv(void *arg, struct udp_pcb *pcb, struct pbuf *p, struct ip_addr *addr, u16_t port)
{

        struct dns_hdr *dns_header;
        //struct ip_addr *answer = arg;
        unsigned char *c; 
        int ans = 0, i;
        
        dns_header = p->payload;
        if(ntohs(dns_header->id) != id) {
          //      sysprintf("wrong id\n");
                goto error;                
        }        
        if(!(ntohs(dns_header->flags) & DNS_RESPONSE)) {
            //    sysprintf("not response\n");
                goto error;
        }
        if(ntohs(dns_header->flags) & 0xf) {	
              //  sysprintf("have error\n");
                goto error; // error
        }
        if((ans = ntohs(dns_header->tot_ans)) == 0) {
                //sysprintf("no response\n");
                goto error; //no ans        
        }
        // skip query 
        c = (unsigned char *)p->payload + sizeof(struct dns_hdr);
        while(*c++ != 0);
        c += 4;  // skip type & class
        
        while(ans-- > 0) {
                if((*c & 0xc0) == 0xc0) {
        	            c += 2;
                }else{
                        while(*c++ != 0);
                }
                if(*c != 0 || *(c + 1) != 1 || *(c + 2) != 0 || *(c + 3) != 1) {
                        c += 8; //type, class & TTL
                        i = *(c) << 8 | *(c + 1); // length
                        c += (i + 2); 
                        continue;
                }
                // Just in case...
                if((unsigned int)c > (unsigned int)((unsigned char *)p->payload + p->tot_len))
                        goto error;
                c += 10; // type, class, TTL, length
                IP4_ADDR(&answer, *(c), *(c + 1), *(c + 2), *(c + 3));
        
                pbuf_free(p);
                flag = 1;
                return;
        }
error:
        pbuf_free(p);
        flag = -1;
        return;

}
#endif